#include <cstring>

#include "common.h"
#include "misc.h"
#include "ndarray_view.h"
#include "tinyndarray.h"
namespace tinyndarray {

int ClipOp(int value, int min_val, int max_val);

Dtype NdArrayView::dtype() const {
    return m_substance->dtype;
}

size_t NdArrayView::size() const {
    size_t s = 1;
    for (auto dim : m_shape) s *= dim;
    return s;
}

NdArray NdArrayView::to_ndarray() const {
    const Dtype dtype = m_substance->dtype;
    const size_t elem_size = get_type_size(dtype);
    NdArray result(dtype, m_shape);

    // 源数据指针（带视图偏移）
    const char* src_base = static_cast<const char*>(m_substance->data.get());
    const char* src = src_base + m_offset;

    // 目标数据指针
    char* dst = static_cast<char*>(result.data_ptr());

    // 总元素数
    const size_t total_elements = size();
    const size_t total_bytes = total_elements * elem_size;

    // 检查视图是否是连续的
    if (is_contiguous()) {
        // 连续内存：直接复制整个块
        memcpy(dst, src, total_bytes);
        return result;
    }

    // 非连续内存：逐元素复制
    RunParallel(total_elements, [&](size_t idx) {
        // 将线性索引转换为多维索引
        std::vector<size_t> indices(m_shape.size());
        size_t remaining = idx;
        for (int i = m_shape.size() - 1; i >= 0; --i) {
            indices[i] = remaining % m_shape[i];
            remaining /= m_shape[i];
        }

        // 计算源元素偏移
        size_t src_offset = 0;
        for (size_t i = 0; i < indices.size(); ++i) {
            src_offset += indices[i] * m_strides[i];
        }

        // 计算目标元素偏移（连续内存）
        const size_t dst_offset = idx * elem_size;

        // 复制元素
        memcpy(dst + dst_offset, src + src_offset, elem_size);
    });

    return result;
}

// 基本切片方法
NdArrayView NdArrayView::slice(const SliceIndex& slice_index) const {
    const size_t ndim = m_shape.size();

    if (slice_index.size() != ndim) {
        throw std::runtime_error("Slice index dimension mismatch");
    }

    // 计算新视图的偏移、形状和步长
    size_t new_offset = m_offset;
    Shape new_shape;
    std::vector<size_t> new_strides = m_strides;  // 步长不变

    for (size_t i = 0; i < ndim; i++) {
        const auto& si = slice_index[i];
        int s = (si.first >= 0) ? si.first : static_cast<int>(m_shape[i]) + si.first;
        int e = (si.second >= 0) ? si.second : static_cast<int>(m_shape[i]) + si.second;

        // 裁剪到有效范围
        s = ClipOp(s, 0, static_cast<int>(m_shape[i]));
        e = ClipOp(e, s + 1, static_cast<int>(m_shape[i]) + 1);

        // 计算维度大小
        const size_t dim_size = std::max(e - s, 0);
        new_shape.push_back(dim_size);

        // 更新偏移量
        new_offset += static_cast<size_t>(s) * m_strides[i];
    }

    return NdArrayView(m_substance, new_offset, new_shape, new_strides);
}

// 便捷切片接口
template <typename... I>
NdArrayView NdArrayView::slice(std::initializer_list<I>... slice_index) const {
    return slice(SliceIndex{CvtToSliceIndexItem(slice_index)...});
}

NdArrayView NdArrayView::operator()(const SliceIndex& slice_index) const {
    return slice(slice_index);
}
template <typename... I>
NdArrayView NdArrayView::operator()(std::initializer_list<I>... slice_index) const {
    return slice(slice_index...);
}

template NdArrayView NdArrayView::slice(IList, IList) const;
template NdArrayView NdArrayView::slice(IList, IList, IList) const;
template NdArrayView NdArrayView::slice(IList, IList, IList, IList) const;
template NdArrayView NdArrayView::slice(IList, IList, IList, IList, IList) const;
template NdArrayView NdArrayView::slice(IList, IList, IList, IList, IList, IList) const;

template NdArrayView NdArrayView::operator()(IList, IList) const;
template NdArrayView NdArrayView::operator()(IList, IList, IList) const;
template NdArrayView NdArrayView::operator()(IList, IList, IList, IList) const;
template NdArrayView NdArrayView::operator()(IList, IList, IList, IList, IList) const;
template NdArrayView NdArrayView::operator()(IList, IList, IList, IList, IList, IList) const;

bool NdArrayView::is_contiguous() const {
    size_t expected_stride = get_type_size(dtype());
    for (int i = m_shape.size() - 1; i >= 0; --i) {
        if (m_strides[i] != expected_stride) {
            return false;
        }
        expected_stride *= m_shape[i];
    }
    return true;
}

// 通用赋值实现
template <typename Source>
NdArrayView& NdArrayView::assign_from(const Source& src) {
    // 检查数据类型是否匹配
    if (dtype() != src.dtype()) {
        throw std::runtime_error("Data type mismatch in assignment");
    }

    // 检查形状是否匹配
    if (shape() != src.shape()) {
        throw std::runtime_error("Shape mismatch in assignment");
    }

    const size_t elem_size = get_type_size(dtype());
    char* dst_base = static_cast<char*>(m_substance->data.get()) + m_offset;
    const char* src_base = get_source_base(src);

    // 检查是否连续内存且布局相同
    if (is_contiguous() && strides() == src.strides()) {
        // 直接内存复制
        const size_t total_bytes = size() * elem_size;
        memcpy(dst_base, src_base, total_bytes);
        return *this;
    }

    // 非连续内存：逐元素复制
    RunParallel(size(), [&](size_t idx) {
        // 将线性索引转换为多维索引
        std::vector<size_t> indices(m_shape.size());
        size_t remaining = idx;
        for (int i = m_shape.size() - 1; i >= 0; --i) {
            indices[i] = remaining % m_shape[i];
            remaining /= m_shape[i];
        }

        // 计算目标元素偏移
        size_t dst_offset = 0;
        for (size_t i = 0; i < indices.size(); ++i) {
            dst_offset += indices[i] * m_strides[i];
        }

        // 计算源元素偏移
        size_t src_offset = calculate_source_offset(src, indices);

        // 复制元素
        memcpy(dst_base + dst_offset, src_base + src_offset, elem_size);
    });

    return *this;
}

template <typename Source>
const char* NdArrayView::get_source_base(const Source& src) const {
    if constexpr (std::is_same_v<Source, NdArray>) {
        return static_cast<const char*>(src.data_ptr());
    } else {
        return static_cast<const char*>(src.m_substance->data.get()) + src.offset();
    }
}

// 计算源数据中的偏移
template <typename Source>
size_t NdArrayView::calculate_source_offset(const Source& src, const std::vector<size_t>& indices) const {
    if constexpr (std::is_same_v<Source, NdArray>) {
        // NdArray 是连续存储的
        size_t linear_index = 0;
        size_t stride = 1;
        for (int i = indices.size() - 1; i >= 0; --i) {
            linear_index += indices[i] * stride;
            stride *= src.shape()[i];
        }
        return linear_index * get_type_size(src.dtype());
    } else {
        // NdArrayView 可能有非连续布局
        size_t offset = 0;
        for (size_t i = 0; i < indices.size(); ++i) {
            offset += indices[i] * src.strides()[i];
        }
        return offset;
    }
}

// 从 NdArray 赋值
NdArrayView& NdArrayView::operator=(const NdArray& src) {
    return assign_from(src);
}

// 从另一个 NdArrayView 赋值
NdArrayView& NdArrayView::operator=(const NdArrayView& src) {
    return assign_from(src);
}

ElementProxy NdArrayView::operator[](size_t index) {
    return (*this)[flat_index_to_multi(index)];
}

// 一维索引（const）
ConstElementProxy NdArrayView::operator[](size_t index) const {
    return (*this)[flat_index_to_multi(index)];
}

// 多维索引（非const）
ElementProxy NdArrayView::operator[](const Index& index) {
    size_t byte_offset = calculate_offset(index, m_shape, m_strides);
    char* data_ptr = static_cast<char*>(m_substance->data.get());
    return ElementProxy(data_ptr + m_offset + byte_offset, dtype());
}

// 多维索引（const）
ConstElementProxy NdArrayView::operator[](const Index& index) const {
    size_t byte_offset = calculate_offset(index, m_shape, m_strides);
    const char* data_ptr = static_cast<const char*>(m_substance->data.get());
    return ConstElementProxy(data_ptr + m_offset + byte_offset, dtype());
}

// 辅助方法：将一维索引转换为多维索引
Index NdArrayView::flat_index_to_multi(size_t index) const {
    if (index >= size()) {
        throw std::out_of_range("Index out of range");
    }

    Index indices;
    size_t remaining = index;
    for (int i = m_shape.size() - 1; i >= 0; --i) {
        indices.insert(indices.begin(), remaining % m_shape[i]);
        remaining /= m_shape[i];
    }
    return indices;
}

std::ostream& operator<<(std::ostream& os, const NdArrayView& arr) {
    os << arr.to_ndarray();
    return os;
}

NdArray operator+(const NdArrayView& lhs, const NdArray& rhs) {
    return lhs.to_ndarray() + rhs;
}
NdArray operator+(const NdArray& lhs, const NdArrayView& rhs) {
    return lhs + rhs.to_ndarray();
}
NdArray operator+(const NdArrayView& lhs, const NdArrayView& rhs) {
    return lhs.to_ndarray() + rhs.to_ndarray();
}

NdArray operator+(const NdArrayView& lhs, const float& rhs) {
    return lhs.to_ndarray() + rhs;
}
}  // namespace tinyndarray